package com.aiwiown.snackmq.broker;

import com.aiwiown.snackmq.broker.config.BrokerConfig;
import com.aiwiown.snackmq.broker.transaction.TransactionCheckerService;
import com.aiwiown.snackmq.cluster.metadata.MetadataManager;
import com.aiwiown.snackmq.cluster.raft.SnackMQStateMachine;
import com.aiwiown.snackmq.cluster.raft.manager.SnackMQRaftNodeManager;
import com.aiwiown.snackmq.cluster.timer.TimerScannerService;
import com.aiwiown.snackmq.common.auth.Action;
import com.aiwiown.snackmq.common.auth.Permission;
import com.aiwiown.snackmq.common.auth.User;
import com.aiwiown.snackmq.common.cluster.BrokerId;
import com.aiwiown.snackmq.common.cluster.MetadataChangeEvent;
import com.aiwiown.snackmq.common.protocol.TopicPartition;
import com.aiwiown.snackmq.management.BrokerMetrics;
import com.aiwiown.snackmq.management.scaling.AutoScalerService;
import com.aiwiown.snackmq.network.connection.ConnectionManager;
import com.aiwiown.snackmq.network.server.NettyServer;
import com.aiwiown.snackmq.storage.partition.LogPartition;
import com.aiwiown.snackmq.storage.service.PartitionedStorageService;
import com.aiwiown.snackmq.storage.service.ZeroCopyPartitionedStorageService;
import com.aiwiown.snackmq.storage.timer.TimerMessageStore;
import com.aiwiown.snackmq.storage.transaction.TransactionLog;
import com.alipay.sofa.jraft.rpc.RpcServer;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

/**
 * 【已重构】BrokerServer 类。
 * 1. 默认凭据管理已从轮询模式重构为事件驱动模式。
 * 2. 组件初始化逻辑已移至 BrokerComponentFactory，构造函数仅负责组装。
 * 3. shutdown() 方法的关闭顺序已优化。
 */
@Slf4j
public class BrokerServer {
    static final int wheelSizeInSeconds = 86400;
    private final NettyServer nettyServer;
    private final PartitionedStorageService storageService;
    private final ZeroCopyPartitionedStorageService zeroCopyStorageService;
    private final Map<TopicPartition, LogPartition> activePartitions = new ConcurrentHashMap<>();
    private final Object partitionManagementLock = new Object();

    @Getter
    private final SnackMQRaftNodeManager dataRaftNodeManager;
    @Getter
    private final MetadataManager metadataManager;
    @Getter
    private final BrokerConfig brokerConfig;
    private final BrokerMetrics brokerMetrics;
    private final AtomicBoolean started = new AtomicBoolean(false);
    private final RpcServer raftRpcServer;
    private final AutoScalerService autoScalerService;
    private final BrokerId selfId;
    private final TimerMessageStore timerMessageStore;
    private final TimerScannerService timerScannerService;
    private final TransactionLog transactionLog;
    private final TransactionCheckerService transactionCheckerService;
    // 【新增】对SnackMQStateMachine的引用
    private final SnackMQStateMachine stateMachine;


    public BrokerServer(BrokerConfig brokerConfig) {
        this.brokerConfig = brokerConfig;
        this.selfId = new BrokerId(brokerConfig.getHost(), brokerConfig.getServerPort());

        // 使用工厂创建所有组件，将“创建”与“组装”分离
        BrokerComponentFactory factory = new BrokerComponentFactory(brokerConfig, activePartitions, selfId);

        this.raftRpcServer = factory.createRaftRpcServer();
        this.metadataManager = factory.createMetadataManager(this.raftRpcServer);
        this.storageService = factory.createStorageService();
        this.zeroCopyStorageService = factory.createZeroCopyStorageService();
        this.timerMessageStore = factory.createTimerMessageStore();
        this.transactionLog = factory.createTransactionLog();

        SnackMQStateMachine stateMachine = factory.createSnackMQStateMachine(this.storageService, this.metadataManager, this.timerMessageStore);
        this.dataRaftNodeManager = factory.createDataRaftNodeManager(stateMachine, this.raftRpcServer);
        // 【新增】保存SnackMQStateMachine的引用
        this.stateMachine = stateMachine;

        this.timerScannerService = factory.createTimerScannerService(this.timerMessageStore, this.storageService, this.dataRaftNodeManager);
        this.transactionCheckerService = factory.createTransactionCheckerService(this.transactionLog);

        this.brokerMetrics = new BrokerMetrics();
        this.autoScalerService = factory.createAutoScalerService(this.brokerMetrics);

        this.nettyServer = factory.createNettyServer(
                this.metadataManager, this.dataRaftNodeManager, this.storageService,
                this.zeroCopyStorageService, stateMachine, this.brokerMetrics,
                this.autoScalerService, this.transactionLog
        );

        // 组件组装：注册监听器和指标
        this.metadataManager.addMetadataListener(this::onMetadataChanged);
        this.brokerMetrics.registerConnectionMetrics(ConnectionManager.getInstance());
    }

    public void start() throws IOException {
        if (!started.compareAndSet(false, true)) {
            log.warn("BrokerServer is already started.");
            return;
        }

        log.info("Starting SnackMQ Broker...");
        try {
            // 启动顺序：底层服务 -> Raft -> 上层应用服务
            this.raftRpcServer.init(null);
            this.storageService.start();
            if (this.zeroCopyStorageService != null) {
                this.zeroCopyStorageService.start();
            }

            if (!this.metadataManager.start(false)) {
                throw new IOException("Failed to start metadata Raft node.");
            }
            if (!this.dataRaftNodeManager.start(false)) {
                throw new IOException("Failed to start data Raft node.");
            }

            this.timerScannerService.start();
            this.transactionCheckerService.start();
            this.autoScalerService.start();
            this.nettyServer.start();
            this.started.set(true);
            log.info("SnackMQ Broker started successfully on port {}. Raft on port {}.", 
                    this.brokerConfig.getServerPort(), this.brokerConfig.getRaftPort());
            
            // 【新增】智能恢复结束机制，监控Raft日志重放状态
            new Thread(() -> {
                try {
                    // 等待至少10秒，确保基本初始化完成
                    Thread.sleep(10000);
                    
                    // 然后每2秒检查一次Raft日志重放状态
                    while (true) {
                        if (this.stateMachine.isRaftLogReplayComplete()) {
                            log.info("检测到Raft日志重放完成: {}", this.stateMachine.getRaftLogReplayStatus());
                            this.stateMachine.endRecovery();
                            log.info("Broker 启动完成，已启用消息写入");
                            break;
                        }
                        
                        log.debug("Raft日志重放仍在进行中: {}", this.stateMachine.getRaftLogReplayStatus());
                        Thread.sleep(2000);
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    log.warn("等待恢复完成的线程被中断", e);
                    // 被中断时也要结束恢复状态
                    this.stateMachine.endRecovery();
                }
            }, "recovery-completion-thread").start();
            Runtime.getRuntime().addShutdownHook(new Thread(this::shutdown));
        } catch (Exception e) {
            log.error("Failed to start SnackMQ Broker. Shutting down...", e);
            shutdown();
            throw new IOException("Broker startup failed", e);
        }
    }

    // setupDefaultCredentialsIfNeeded() 方法保持不变 ...
    private void setupDefaultCredentialsIfNeeded() {
        // 只在认证功能启用时执行
        if (!brokerConfig.isAuthEnabled()) {
            return;
        }

        // 确保只有元数据 Leader 才执行此操作
        if (metadataManager.isLeader()) {
            try {
                final String adminUser = brokerConfig.getDefaultAdminUser();
                final String adminRole = "admin-role";

                // --- 【最终修复】将用户创建和权限设置解耦，以处理部分失败的场景 ---

                // 1. 检查并创建用户（如果需要）
                if (metadataManager.getUser(adminUser) == null) {
                    log.info("This node is the leader and default admin user '{}' was not found. Attempting to create it.", adminUser);
                    User user = new User(adminUser, brokerConfig.getDefaultAdminPassword(), Collections.singletonList(adminRole));
                    metadataManager.createUser(user, status -> {
                        if (status.isOk()) {
                            log.info("Successfully proposed creation of default admin user.");
                        } else {
                            log.error("Failed to propose creation of default admin user: {}", status.getErrorMsg());
                        }
                    });
                }

                // 2. 独立地检查并设置角色权限（如果需要）
                // 即使 admin 用户已存在，如果 Broker 在设置权限之前崩溃，此检查将确保权限被正确设置。
                List<Permission> requiredPermissions = getRequiredAdminPermissions();
                List<Permission> currentPermissions = metadataManager.getPermissionsForRole(adminRole);

                // 【最终修复】检查当前权限是否与必需权限完全一致，如果不一致则强制覆盖。
                // 这可以修复因版本迭代或部分失败导致的权限不完整问题。
                if (currentPermissions == null || !new HashSet<>(currentPermissions).equals(new HashSet<>(requiredPermissions))) {
                    log.info("This node is the leader and permissions for '{}' are missing or incorrect. Attempting to set/overwrite them.", adminRole);

                    metadataManager.setRolePermissions(adminRole, requiredPermissions, permStatus -> {
                        if (permStatus.isOk()) {
                            log.info("Successfully proposed setting default permissions for '{}'.", adminRole);
                        } else {
                            log.error("Failed to propose setting permissions for '{}': {}", adminRole, permStatus.getErrorMsg());
                        }
                    });
                }
            } catch (Exception e) {
                log.error("An unexpected error occurred during default credential setup.", e);
            }
        }
    }

    private List<Permission> getRequiredAdminPermissions() {
        List<Permission> adminPermissions = new ArrayList<>();
        adminPermissions.add(new Permission("CLUSTER", "*", Action.MANAGE));
        adminPermissions.add(new Permission("TOPIC", "*", Action.MANAGE));
        adminPermissions.add(new Permission("TOPIC", "*", Action.PRODUCE));
        adminPermissions.add(new Permission("TOPIC", "*", Action.CONSUME));
        adminPermissions.add(new Permission("EXCHANGE", "*", Action.PRODUCE));
        adminPermissions.add(new Permission("CONSUMER_GROUP", "*", Action.MANAGE));
        return adminPermissions;
    }

    public void shutdown() {
        if (!started.compareAndSet(true, false)) {
            return;
        }
        log.info("Shutting down SnackMQ Broker...");

        // 1. 【优化】关闭所有上层应用服务，停止接受新请求和后台任务
        if (this.nettyServer != null) this.nettyServer.shutdown();
        if (this.autoScalerService != null) this.autoScalerService.close();
        if (this.timerScannerService != null) this.timerScannerService.shutdown();
        if (this.transactionCheckerService != null) this.transactionCheckerService.shutdown();

        // 2. 关闭 Raft 节点，停止集群协调
        if (this.dataRaftNodeManager != null) this.dataRaftNodeManager.shutdown();
        if (this.metadataManager != null) this.metadataManager.shutdown();

        // 3. 关闭 RPC 服务
        if (this.raftRpcServer != null) this.raftRpcServer.shutdown();

        // 4. 关闭所有活动的分区
        synchronized (partitionManagementLock) {
            activePartitions.values().forEach(partition -> {
                try {
                    partition.close();
                } catch (IOException e) {
                    log.error("Error closing partition during broker shutdown.", e);
                }
            });
            activePartitions.clear();
        }

        // 5. 【优化】最后关闭所有底层存储服务
        if (this.timerMessageStore != null) {
            try {
                this.timerMessageStore.close();
            } catch (IOException e) {
                log.error("Error closing timer message store.", e);
            }
        }
        if (this.zeroCopyStorageService != null) {
            try {
                this.zeroCopyStorageService.close();
            } catch (IOException e) {
                log.error("Error closing zero-copy storage service.", e);
            }
        }
        if (this.transactionLog != null) {
            try {
                this.transactionLog.close();
            } catch (IOException e) {
                log.error("Error closing transaction log.", e);
            }
        }
        if (this.storageService != null) {
            try {
                this.storageService.close();
            } catch (IOException e) {
                log.error("Error closing storage service.", e);
            }
        }
        log.info("SnackMQ Broker shutdown complete.");
    }

    // onMetadataChanged, updateLocalPartitions, isStarted, main 等方法保持不变 ...
    @SuppressWarnings("unchecked")
    private void onMetadataChanged(MetadataChangeEvent event) {
        // --- 【优化】: 增加对 LEADER_ELECTED 事件的处理 ---
        if (event.getType() == MetadataChangeEvent.EventType.METADATA_LEADER_ELECTED) {
            log.info("Broker {} received METADATA_LEADER_ELECTED event, attempting to set up default credentials.", selfId);
            setupDefaultCredentialsIfNeeded();
        } else if (event.getType() == MetadataChangeEvent.EventType.PARTITION_ASSIGNMENT_CHANGED) {
            log.info("Broker {} received PARTITION_ASSIGNMENT_CHANGED event.", selfId);
            Map<TopicPartition, List<BrokerId>> newAssignment =
                    (Map<TopicPartition, List<BrokerId>>) event.getContext().get("newAssignment");

            if (newAssignment != null) {
                updateLocalPartitions(newAssignment);
            }
        }
    }

    private void updateLocalPartitions(Map<TopicPartition, List<BrokerId>> newAssignment) {
        synchronized (partitionManagementLock) {
            Set<TopicPartition> newLocalPartitions = newAssignment.entrySet().stream()
                    .filter(entry -> entry.getValue().contains(selfId))
                    .map(Map.Entry::getKey)
                    .collect(Collectors.toSet());
            Set<TopicPartition> currentLocalPartitions = new HashSet<>(this.activePartitions.keySet());
            Set<TopicPartition> partitionsToCreate = new HashSet<>(newLocalPartitions);
            partitionsToCreate.removeAll(currentLocalPartitions);
            Set<TopicPartition> partitionsToRemove = new HashSet<>(currentLocalPartitions);
            partitionsToRemove.removeAll(newLocalPartitions);

            partitionsToRemove.forEach(tp -> {
                log.info("Partition {} is no longer assigned to Broker {}. Closing and removing.", tp, selfId);
                LogPartition partition = this.activePartitions.remove(tp);
                if (partition != null) {
                    try {
                        partition.close();
                    } catch (IOException e) {
                        log.error("Failed to close partition {} during re-assignment.", tp, e);
                    }
                }
            });

            partitionsToCreate.forEach(tp -> {
                log.info("Partition {} is now assigned to Broker {}. Creating.", tp, selfId);
                try {
                    if (!this.activePartitions.containsKey(tp)) {
                        LogPartition partition = new LogPartition(
                                brokerConfig.getStorePath(),
                                tp.getTopic(),
                                tp.getPartitionId(),
                                brokerConfig.getSegmentSize(),
                                brokerConfig.getIndexInterval(),
                                brokerConfig.isForceFullLogRecovery()
                        );
                        // 先设置storageService
                        if (this.storageService instanceof PartitionedStorageService) {
                            partition.setStorageService((PartitionedStorageService) this.storageService);
                        }
                        // 再初始化和恢复
                        partition.initializeAndRecover();
                        this.activePartitions.put(tp, partition);
                    }
                } catch (IOException e) {
                    log.error("Failed to create partition {}.", tp, e);
                }
            });
        }
    }

    public boolean isStarted() {
        return started.get();
    }

    public static void main(String[] args) {
        BrokerConfig config = new BrokerConfig();
        // --- 【核心修复】: 在启动 Broker 前启用认证 ---
        config.setAuthEnabled(true);
        BrokerServer server = new BrokerServer(config);
        try {
            server.start();
        } catch (IOException e) {
            log.error("Failed to start broker server", e);
            System.exit(1);
        }
    }
}